{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SPU Basics\n",
    "\n",
    ">The following codes are demos only. It's **NOT for production** due to system security concerns, please **DO NOT** use it directly in production.\n",
    "\n",
    "SPU devices are responsible for performing MPC computation in SecretFlow.\n",
    "\n",
    "This tutorial would help you:\n",
    "\n",
    "- be familiar with SPU device and SPU Object\n",
    "- learn how to transfer a Python Object / PYU Object from/to SPU Object.\n",
    "- run MPC computation with SPU device.\n",
    "\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create an SPU Device\n",
    "\n",
    "### Create SecretFlow Parties\n",
    "\n",
    "Parties are basic nodes in SecretFlow nodes. We are going to create four parties - **alice**, **bob**, **carol** and **dave**.\n",
    "\n",
    "Based on four parties, we will set up three devices:\n",
    "\n",
    "- a PYU device based on *alice*\n",
    "- a PYU device based on *dave*\n",
    "- an SPU device based on *alice*, *bob* and *carol*\n",
    "\n",
    "<img alt=\"spu_basics_devices.png\" src=\"resources/spu_basics_devices.png\">\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-06-17 18:30:38,459\tINFO worker.py:1538 -- Started a local Ray instance.\n"
     ]
    }
   ],
   "source": [
    "import secretflow as sf\n",
    "\n",
    "# Check the version of your SecretFlow\n",
    "print('The version of SecretFlow: {}'.format(sf.__version__))\n",
    "\n",
    "# In case you have a running secretflow runtime already.\n",
    "sf.shutdown()\n",
    "\n",
    "sf.init(['alice', 'bob', 'carol', 'dave'], address='local')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create a 3PC ABY3 SPU device"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After that, let's create an SPU device with [ABY3](https://eprint.iacr.org/2018/403.pdf) protocol.\n",
    "\n",
    "`sf.utils.testing.cluster_def` is a helper method to create a config by finding unused ports."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'nodes': [{'party': 'alice', 'address': '127.0.0.1:49613'},\n",
       "  {'party': 'bob', 'address': '127.0.0.1:52053'},\n",
       "  {'party': 'carol', 'address': '127.0.0.1:25589'}],\n",
       " 'runtime_config': {'protocol': 3, 'field': 3}}"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "aby3_config = sf.utils.testing.cluster_def(parties=['alice', 'bob', 'carol'])\n",
    "\n",
    "aby3_config"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then let's use *aby3_config* to create an SPU device and check its cluster_def."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'nodes': [{'party': 'alice', 'address': '127.0.0.1:49613'},\n",
       "  {'party': 'bob', 'address': '127.0.0.1:52053'},\n",
       "  {'party': 'carol', 'address': '127.0.0.1:25589'}],\n",
       " 'runtime_config': {'protocol': 3, 'field': 3}}"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "spu_device = sf.SPU(aby3_config)\n",
    "\n",
    "spu_device.cluster_def"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lastly, let's create two PYU devices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "alice, dave = sf.PYU('alice'), sf.PYU('dave')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pass Values to SPU device\n",
    "\n",
    "Before talking about computation with SPU device, let's understand how to pass a Python object or a PYUObject to SPU device."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### SPUObject\n",
    "\n",
    "A Python object or a PYUObject could be transferred into an SPUObject and secret-shared by SPU nodes.\n",
    "\n",
    "**sf.device.SPUIO** is the helper class to do the job. You don't need to call this method in your code. We just use it to demonstrate the structure of **SPUObjects** and everything happens for you.\n",
    "\n",
    "Each SPUObject has two fields:\n",
    "\n",
    "- meta: The structure of the origin object.\n",
    "- shares: The secret sharing of the origin object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:jax._src.xla_bridge:Unable to initialize backend 'cuda': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'\n",
      "INFO:jax._src.xla_bridge:Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'\n",
      "INFO:jax._src.xla_bridge:Unable to initialize backend 'tpu': INVALID_ARGUMENT: TpuPlatform is not available.\n",
      "INFO:jax._src.xla_bridge:Unable to initialize backend 'plugin': xla_extension has no attributes named get_plugin_device_client. Compile TensorFlow with //tensorflow/compiler/xla/python:enable_plugin_device set to true (defaults to false) to enable this.\n",
      "WARNING:jax._src.xla_bridge:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n"
     ]
    }
   ],
   "source": [
    "spu_io = sf.device.SPUIO(spu_device.conf, spu_device.world_size)\n",
    "\n",
    "bank_account = [{'id': 12345, 'deposit': 1000.25}, {'id': 12345, 'deposit': 100000.25}]\n",
    "\n",
    "import spu\n",
    "\n",
    "meta, io_info, *shares = spu_io.make_shares(bank_account, spu.Visibility.VIS_SECRET)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's check meta first."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'deposit': SPUValueMeta(shape=(), dtype=dtype('float32'), vtype=1, protocol=3, field=3, fxp_fraction_bits=0),\n",
       "  'id': SPUValueMeta(shape=(), dtype=dtype('int32'), vtype=1, protocol=3, field=3, fxp_fraction_bits=0)},\n",
       " {'deposit': SPUValueMeta(shape=(), dtype=dtype('float32'), vtype=1, protocol=3, field=3, fxp_fraction_bits=0),\n",
       "  'id': SPUValueMeta(shape=(), dtype=dtype('int32'), vtype=1, protocol=3, field=3, fxp_fraction_bits=0)}]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "meta"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I guess you could find meta preserves the structure of origin data and replaces the digits/arrays with **SPUValueMeta**:\n",
    "\n",
    "- data_type, indicates whether the value is integer or fixed points.\n",
    "- visibility, indicates whether the value is a secret or a public content. \n",
    "- storage_type, indicates attributes of value, e.g. MPC protocol(ABY3 in our case), field size(128 bits in our case), etc\n",
    "\n",
    "Then let's check shares of bank_account_spu. Since we are passing data to a 3PC SPU device. We would have three pieces of shares,\n",
    "and we are going to check the first piece."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'deposit': b'\\x08\\n\\x10\\x01\"\\x10aby3.AShr<FM128>* \\xcd\\xbd\\xed#\\x06\\x04\\x0f\\xebJ\\xdc\\xdf\\x1b\\xacUe\\xdc\\xbe\\'\\x94\\xbb\\xf8?\\xa9-\\x99\\xc8TzM\\xf3\\xe4\\xaf',\n",
       "  'id': b'\\x08\\x06\\x10\\x01\"\\x10aby3.AShr<FM128>* \\xf0\\x8b\\xaa\\xc4\\xe5V\\x8a^\\xffq>\\xee\\x08\\x85\\xa6\\x87\\x82C\\xb6\\xbf|_\\xff\\x18\\xfb\\xb7\\xe3`\\x86\\xea\\xc9\\x1a'},\n",
       " {'deposit': b'\\x08\\n\\x10\\x01\"\\x10aby3.AShr<FM128>* \\xbaB\\x18\\xa6\\x84\\x9eW\\xa3\\xe8\\x18\\xc6\\x81\\xc7\\x1dp\\'\\x03\\xb4\\xa7\\xa6\\x9e\\x0eF\\xfan\\x81\\xd33,\\xcd\\x05X',\n",
       "  'id': b'\\x08\\x06\\x10\\x01\"\\x10aby3.AShr<FM128>* xj\\xde\\x12\\xa9\\x82\\xdfi\\xaahZ\\x16\\r\\xdeH\\x15$\\x17\\xce\\x05\\x8f\\x9b\\x9f\\xc5\\x81d\\x94!\\xab\\x983\\xaf'}]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "assert len(shares) == 12\n",
    "\n",
    "shares[0]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You should find a piece of shares of SPU Object is very similar to meta and origin data. It still preserves the structure of origin data while digits are replaced by encoded secret (try to guess the origin data if you would like to).\n",
    "\n",
    "Well, let's reconstruct the origin Python object from SPU object. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'deposit': array(1000.25, dtype=float32), 'id': array(12345, dtype=int32)},\n",
       " {'deposit': array(100000.25, dtype=float32), 'id': array(12345, dtype=int32)}]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bank_account_hat = spu_io.reconstruct(shares, io_info, meta)\n",
    "bank_account_hat"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you compare **bank_account_hat** with origin **bank_account**, you should find all the digits in **bank_account_hat** have become **numpy.array** but values are preseved."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Pass a PYU Object from PYU to SPU\n",
    "\n",
    "First, we create a PYU object with a PYU device."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<secretflow.device.device.pyu.PYUObject at 0x7fd98cd09130>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def debit_amount():\n",
    "    return 10\n",
    "\n",
    "\n",
    "debit_amount_pyu = alice(debit_amount)()\n",
    "debit_amount_pyu"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then let's pass debit_amount_pyu from PYU to SPU. We will get an SPU object as result. Under the hood, **alice** calls **sf.device.SPUIO.make_shares** to get **meta** and **shares** to send to nodes of the spu device."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<secretflow.device.device.spu.SPUObject at 0x7fd817a03c70>"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "debit_amount_spu = debit_amount_pyu.to(spu_device)\n",
    "\n",
    "debit_amount_spu"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's check meta of debit_amount_spu."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ObjectRef(e0dc174c83599034ffffffffffffffffffffffff0100000001000000)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "debit_amount_spu.meta"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Oh no, it's a Ray ObjectRef located at alice part.\n",
    "So how about shares of debit_amount_spu?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[ObjectRef(f4402ec78d3a260750696baee0bc0bb42b40620a0100000001000000),\n",
       " ObjectRef(f91b78d7db9a65936b44b364879d9518bec82ea10100000001000000),\n",
       " ObjectRef(82891771158d68c155ebf101d0aa7682c810dad40100000001000000)]"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "debit_amount_spu.shares_name"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So you get a list of ObjectRef! Since it's located at alice part, we couldn't check the value at host."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But if you are really curious, we could use **sf.reveal** to check the origin value. Be careful to use **sf.reveal** in production! When **sf.reveal** are applied on **SPUObjects**, **sf.device.SPUIO.reconstruct** are called for you."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[2m\u001b[36m(_run pid=102815)\u001b[0m INFO:jax._src.xla_bridge:Unable to initialize backend 'cuda': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'\n",
      "\u001b[2m\u001b[36m(_run pid=102815)\u001b[0m INFO:jax._src.xla_bridge:Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'\n",
      "\u001b[2m\u001b[36m(_run pid=102815)\u001b[0m INFO:jax._src.xla_bridge:Unable to initialize backend 'tpu': INVALID_ARGUMENT: TpuPlatform is not available.\n",
      "\u001b[2m\u001b[36m(_run pid=102815)\u001b[0m INFO:jax._src.xla_bridge:Unable to initialize backend 'plugin': xla_extension has no attributes named get_plugin_device_client. Compile TensorFlow with //tensorflow/compiler/xla/python:enable_plugin_device set to true (defaults to false) to enable this.\n",
      "\u001b[2m\u001b[36m(_run pid=102815)\u001b[0m WARNING:jax._src.xla_bridge:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array(10, dtype=int32)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sf.reveal(debit_amount_spu)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Pass a Python Object from Host to SPU\n",
    "\n",
    "Let's pass a dict from Host to SPU device. \n",
    "\n",
    "> NOTE: I know it looks weird. At this moment, if you want to pass a Python object to SPU device, you have to pass it to a PYU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "bank_account_spu = sf.to(alice, bank_account).to(spu_device)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Summary\n",
    "\n",
    "\n",
    "This is the first part of Data Flow with SPU device, at this moment, you should be aware of the following facts.\n",
    "\n",
    "- A Python Object/PYU Object could be transferred to an SPU Object.\n",
    "- An SPU Object consists of meta and shares.\n",
    "- **sf.to** and **sf.reveal** calls **sf.device.SPUIO** to transfer between SPUObjects and Python objects.\n",
    "- Just converting to SPU Object won't trigger data flow from PYU to SPU. e.g. When you transferred a PYU object to an SPU object. All the field of SPU objects including meta and shares are still located at the PYU device. The shares would only be sent to parties of SPU device when computation do happen. In short, data flow is lazy."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Computation with SPU Device\n",
    "\n",
    "Since we have two SPU objects - *bank_account_spu* and *debit_amount_spu* as inputs.\n",
    "Let's try to do some computation with SPU device.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<secretflow.device.device.spu.SPUObject at 0x7fd98cca88b0>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def deduce_from_account(bank_account, amount):\n",
    "    new_bank_account = []\n",
    "\n",
    "    for account in bank_account:\n",
    "        account['deposit'] = account['deposit'] - amount\n",
    "        new_bank_account.append(account)\n",
    "\n",
    "    return new_bank_account\n",
    "\n",
    "\n",
    "new_bank_account_spu = spu_device(deduce_from_account)(\n",
    "    bank_account_spu, debit_amount_spu\n",
    ")\n",
    "\n",
    "new_bank_account_spu"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*new_bank_account_spu* is also a **SPUObject**. But it's a bit different from *debit_amount_spu*!\n",
    "\n",
    "- *debit_amount_spu* is located at alice, so only alice could check value.\n",
    "- *new_bank_account_spu* is located at spu, each party of spu have a piece of shares. And you couldn't check the value directly without *sf.reveal*.\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Well, but what happened behind computation of SPU device?"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 1: Compile Python(Jax) Code to SPU Executable\n",
    "\n",
    "The Python function (*deduce_from_account* in our case) and metas of all inputs (*bank_account_spu* and *debit_amount_spu*) would be sent to one party of SPU device. Then SPU compiler would be used to compile them to *SPU Executable*.\n",
    "\n",
    "<img alt=\"spu_basics_compiler.png\" src=\"resources/spu_basics_compiler.png\">"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2: Distribute the SPU Executable and Shares to SPU parties.\n",
    "\n",
    "Each party of SPU device would get:\n",
    "\n",
    "- one copy of SPU Executable\n",
    "- one piece of each SPU Object share\n",
    "\n",
    "<img alt=\"spu_basics_distribute.png\" src=\"resources/spu_basics_distribute.png\">"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 3: Run SPU Executable and Assembly SPU Object\n",
    "\n",
    "Then each party of SPU device would execute SPU Executable.\n",
    "\n",
    "In the end, each party of SPU device would own a piece of output SPU Objects and a copy of meta.\n",
    "\n",
    "Then SecretFlow framework would use them to assembly SPU Objects."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get Value from SPU Device\n",
    "\n",
    "But in the end, we need to get value from spu, we couldn't always keep *SPUObject* as secret!\n",
    "\n",
    "Most common way of handling *SPUObject* is pass the secret to one party. This party is not necessarily one of parties consisting of SPU device."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<secretflow.device.device.pyu.PYUObject at 0x7fd98cd754f0>"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_bank_account_pyu = new_bank_account_spu.to(dave)\n",
    "\n",
    "new_bank_account_pyu"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We just pass *new_bank_account_spu* to **pyu**, then it becomes a *PYUObject*! And it's owned by dave.\n",
    "Let's check the value of *new_bank_account_pyu*."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'deposit': array(990.25, dtype=float32), 'id': array(12345, dtype=int32)},\n",
       " {'deposit': array(99990.25, dtype=float32), 'id': array(12345, dtype=int32)}]"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sf.reveal(new_bank_account_pyu)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We could also pass *SPUObject* to host directly. The magic is *sf.reveal*. And again, be careful in production!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'deposit': array(990.25, dtype=float32), 'id': array(12345, dtype=int32)},\n",
       " {'deposit': array(99990.25, dtype=float32), 'id': array(12345, dtype=int32)}]"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sf.reveal(new_bank_account_spu)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Advanced Topic: Use Different MPC Protocol\n",
    "\n",
    "At this moment, SPU device supports multiple MPC protocol besides ABY3. It's easy to use different MPC protocol - just set the proper field in cluster def.\n",
    "\n",
    "For instance, if someone would like to use 2PC protocol - Cheetah,\n",
    "You should prepare another cluster def:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-06-17 18:30:47,897\tINFO worker.py:1538 -- Started a local Ray instance.\n"
     ]
    }
   ],
   "source": [
    "import spu\n",
    "\n",
    "import secretflow as sf\n",
    "\n",
    "# In case you have a running secretflow runtime already.\n",
    "sf.shutdown()\n",
    "\n",
    "sf.init(['alice', 'bob', 'carol', 'dave'], address='local')\n",
    "\n",
    "cheetah_config = sf.utils.testing.cluster_def(\n",
    "    parties=['alice', 'bob'],\n",
    "    runtime_config={\n",
    "        'protocol': spu.spu_pb2.CHEETAH,\n",
    "        'field': spu.spu_pb2.FM64,\n",
    "    },\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then you could create an SPU device with *cheetah_config*."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "spu_device2 = sf.SPU(cheetah_config)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's check the *cluster_def* of spu_device2."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'nodes': [{'party': 'alice', 'address': '127.0.0.1:64555'},\n",
       "  {'party': 'bob', 'address': '127.0.0.1:30243'}],\n",
       " 'runtime_config': {'protocol': 4, 'field': 2}}"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "spu_device2.cluster_def"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We could use *spu_device2* to check famous Yao's Millionaires' problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_carol_assets():\n",
    "    return 1000000\n",
    "\n",
    "\n",
    "def get_dave_assets():\n",
    "    return 1000002\n",
    "\n",
    "\n",
    "carol, dave = sf.PYU('carol'), sf.PYU('dave')\n",
    "\n",
    "carol_assets = carol(get_carol_assets)()\n",
    "dave_assets = dave(get_dave_assets)()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We use *spu_device2* to check if *carol* is richer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[2m\u001b[36m(_run pid=112466)\u001b[0m INFO:jax._src.xla_bridge:Unable to initialize backend 'cuda': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'\n",
      "\u001b[2m\u001b[36m(_run pid=112466)\u001b[0m INFO:jax._src.xla_bridge:Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'\n",
      "\u001b[2m\u001b[36m(_run pid=112466)\u001b[0m INFO:jax._src.xla_bridge:Unable to initialize backend 'tpu': INVALID_ARGUMENT: TpuPlatform is not available.\n",
      "\u001b[2m\u001b[36m(_run pid=112466)\u001b[0m INFO:jax._src.xla_bridge:Unable to initialize backend 'plugin': xla_extension has no attributes named get_plugin_device_client. Compile TensorFlow with //tensorflow/compiler/xla/python:enable_plugin_device set to true (defaults to false) to enable this.\n",
      "\u001b[2m\u001b[36m(_run pid=112466)\u001b[0m WARNING:jax._src.xla_bridge:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n",
      "\u001b[2m\u001b[36m(_run pid=112459)\u001b[0m INFO:jax._src.xla_bridge:Unable to initialize backend 'cuda': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'\n",
      "\u001b[2m\u001b[36m(_run pid=112459)\u001b[0m INFO:jax._src.xla_bridge:Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'\n",
      "\u001b[2m\u001b[36m(_run pid=112459)\u001b[0m INFO:jax._src.xla_bridge:Unable to initialize backend 'tpu': INVALID_ARGUMENT: TpuPlatform is not available.\n",
      "\u001b[2m\u001b[36m(_run pid=112459)\u001b[0m INFO:jax._src.xla_bridge:Unable to initialize backend 'plugin': xla_extension has no attributes named get_plugin_device_client. Compile TensorFlow with //tensorflow/compiler/xla/python:enable_plugin_device set to true (defaults to false) to enable this.\n",
      "\u001b[2m\u001b[36m(_run pid=112459)\u001b[0m WARNING:jax._src.xla_bridge:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array(False)"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def get_winner(carol, dave):\n",
    "    return carol > dave\n",
    "\n",
    "\n",
    "winner = spu_device2(get_winner)(carol_assets, dave_assets)\n",
    "\n",
    "sf.reveal(winner)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Advanced Topic: Multiple Returns from SPU Computation\n",
    "\n",
    "In most cases, we have multiple returns from the function executed by SPU device.\n",
    "\n",
    "For instance,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_multiple_outputs(x, y):\n",
    "    return x + y, x - y"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are multiple options to handle this.\n",
    "\n",
    "### Option 1: Treat All Returns as Single\n",
    "\n",
    "This is the default behavior of SPU. Let's see."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<secretflow.device.device.spu.SPUObject at 0x7fd98cd754c0>"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "single_output = spu_device2(get_multiple_outputs)(carol_assets, dave_assets)\n",
    "\n",
    "single_output"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We could see we only get a single *SPUObject*. Let's reveal it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array(2000002, dtype=int32), array(-2, dtype=int32))"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sf.reveal(single_output)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So single_output itself actually represents a tuple."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Option 2: Decide Return Nums on the Fly\n",
    "\n",
    "We can also instruct SPU to decide return numbers for us."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(<secretflow.device.device.spu.SPUObject at 0x7fd98cce0400>,\n",
       " <secretflow.device.device.spu.SPUObject at 0x7fd98cce0490>)"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from secretflow.device.device.spu import SPUCompilerNumReturnsPolicy\n",
    "\n",
    "multiple_outputs = spu_device2(\n",
    "    get_multiple_outputs, num_returns_policy=SPUCompilerNumReturnsPolicy.FROM_COMPILER\n",
    ")(carol_assets, dave_assets)\n",
    "\n",
    "multiple_outputs"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "let's check two outputs respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2000002\n",
      "-2\n"
     ]
    }
   ],
   "source": [
    "print(sf.reveal(multiple_outputs[0]))\n",
    "print(sf.reveal(multiple_outputs[1]))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Option 3: Decide Return Nums Manually\n",
    "\n",
    "If possible, you could also set return nums manually."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<secretflow.device.device.spu.SPUObject at 0x7fd98cce0a60>,\n",
       " <secretflow.device.device.spu.SPUObject at 0x7fd98cce0af0>]"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "user_multiple_outputs = spu_device2(\n",
    "    get_multiple_outputs,\n",
    "    num_returns_policy=SPUCompilerNumReturnsPolicy.FROM_USER,\n",
    "    user_specified_num_returns=2,\n",
    ")(carol_assets, dave_assets)\n",
    "\n",
    "user_multiple_outputs"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "let's also check two outputs respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2000002\n",
      "-2\n"
     ]
    }
   ],
   "source": [
    "print(sf.reveal(multiple_outputs[0]))\n",
    "print(sf.reveal(multiple_outputs[1]))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's summarize what we have:\n",
    "\n",
    "- Be default, SPU treats all the returns as a single return\n",
    "- Since SPU compiler generates the SPU executable, it can figure out return nums. However, the options results some latency since we have to make compilation blocked.\n",
    "- If you want to avoid latency, we can provide return nums manually. But you have to make sure you provide the right nums, otherwise, the program would complain!"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What's Next?\n",
    "\n",
    "After learning basics of SPU, you may check some advanced tutorials with SPU:\n",
    "\n",
    "- [Logistic Regression with SPU](./lr_with_spu.ipynb)\n",
    "- [Neural Network with SPU](./nn_with_spu.ipynb)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.16"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "de4c945f5346493decaa0ea82289843a7da2415616b96b9f4b104111cc0c19ad"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
